from py2neo import Graph, Node, Relationship
import pandas as pd
import time
from tqdm import tqdm


def import_data_to_neo4j(csv_path, neo4j_uri, neo4j_user, neo4j_password):
    # -------------------------- 1. 读取CSV数据，确认列结构 --------------------------
    try:
        # 读取CSV，指定空值标识（兼容空字符串、NA等）
        data = pd.read_csv(csv_path, na_values=['', 'NA', 'na'], keep_default_na=True)
        print(f"✅ 成功读取数据：共 {len(data)} 行，{len(data.columns)} 列")

        # 打印所有列名（确认列→实体标签的映射关系）
        print("\n📊 CSV列名（对应Neo4j节点标签）：")
        for idx, col in enumerate(data.columns, 1):
            print(f"   {idx}. 列名：{col} → 节点标签：{col}")


    except Exception as e:
        print(f"❌ 读取CSV失败：{str(e)}")
        return

    # -------------------------- 2. 连接Neo4j数据库 --------------------------
    try:
        graph = Graph(neo4j_uri, auth=(neo4j_user, neo4j_password))
        # 测试连接（执行简单查询）
        test_count = graph.run("MATCH (n) RETURN count(n) LIMIT 1").evaluate()
        print(f"\n✅ 成功连接Neo4j：当前数据库节点总数 = {test_count}")

        # 可选：清除现有数据（首次导入用，后续增量导入请注释）
        confirm_clear = input("\n❓ 是否清除现有数据（首次导入建议清除，输入 y/n）：")
        if confirm_clear.lower() == 'y':
            graph.run("MATCH (n) DETACH DELETE n")
            print("✅ 已清除现有数据")

    except Exception as e:
        error_msg = str(e).lower()
        if "auth" in error_msg:
            print(f"❌ Neo4j认证失败：请检查用户名/密码")
        elif "connection" in error_msg:
            print(f"❌ Neo4j连接失败：请检查URI/端口/容器状态")
        else:
            print(f"❌ Neo4j连接错误：{str(e)}")
        return

    # -------------------------- 3. 按列导入实体节点（核心逻辑） --------------------------
    print(f"\n🚀 开始按列导入实体节点（共 {len(data)} 行数据）：")
    start_time = time.time()
    node_count = 0  # 统计总节点数

    # 遍历每一行数据（tqdm显示进度）
    for row_idx, (_, row) in tqdm(enumerate(data.iterrows(), 1), total=len(data), desc="导入进度"):
        # 存储当前行创建的节点（用于后续建立关系）
        current_row_nodes = {}

        # 遍历每一列（按列创建实体节点）
        for col_name in data.columns:
            # 获取当前单元格值（去除前后空格）
            cell_value = str(row[col_name]).strip() if pd.notna(row[col_name]) else None

            # 跳过空值（不创建空节点）
            if not cell_value or cell_value in ['nan', 'none']:
                continue

            # -------------------------- 3.1 定义节点属性（确保唯一性） --------------------------
            # 节点标签：直接用列名（如“致甜成分实验方法”列→标签“致甜成分实验方法”）
            node_label = col_name
            # 唯一ID：列名_行索引_值哈希（避免同值不同行的节点重复，如“致甜成分_5_麦芽酚”）
            unique_id = f"{col_name}_{row_idx}_{hash(cell_value)}"
            # 节点属性：至少包含“唯一ID”“值”“来源行号”（便于追溯数据）
            node_properties = {
                "id": unique_id,
                "name": cell_value,
                "source_row": row_idx,  # 记录数据来源行号，便于调试
                "create_time": time.strftime("%Y-%m-%d %H:%M:%S")
            }

            # -------------------------- 3.2 创建/合并节点（避免重复） --------------------------
            # 用“标签+id”合并节点（确保唯一，即使值相同，不同行/列也不会误合并）
            node = Node(node_label, **node_properties)
            graph.merge(node, node_label, "id")  # 关键：merge依据“标签+id”
            node_count += 1

            # -------------------------- 3.3 记录核心节点（用于后续建立关系） --------------------------
            node_label = col_name  # 标签直接用列名（如“致甜成分实验方法”列→标签“致甜成分实验方法”）
            current_row_nodes[node_label] = node  # 存储：标签→节点对象（所有列的节点都记录）

        # -------------------------- 4. 建立核心业务关系（关键步骤） --------------------------
        # -------------------------- 4. 自定义所有列之间的关系 --------------------------
        if len(current_row_nodes) >= 2:
            # 示例1：甜味剂类型 → 研究论文 → 论文名称
            if "甜味剂类型" in current_row_nodes and "论文名称" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["甜味剂类型"],  # 源节点
                        "研究论文",  # 关系类型
                        current_row_nodes["论文名称"]  # 目标节点
                    ),
                    "研究论文"  # 关系唯一标识依据

                )

            # 示例2：论文名称 → 致甜物质 → 致甜成分名称
            if "论文名称" in current_row_nodes and "致甜成分名称" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["论文名称"],
                        "致甜物质",
                        current_row_nodes["致甜成分名称"]
                    ),
                    "致甜物质"
                )

            # 示例3：致甜成分名称 → 应用于 → 致甜成分的应用
            if "致甜成分名称" in current_row_nodes and "致甜成分的应用" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "应用于",
                        current_row_nodes["致甜成分的应用"]
                    ),
                    "应用于"
                )

            if "致甜成分名称" in current_row_nodes and "卷烟余味" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "余味",
                        current_row_nodes["卷烟余味"]
                    ),
                    "余味"
                )

            if "致甜成分名称" in current_row_nodes and "烟叶类型" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "烟叶类别",
                        current_row_nodes["烟叶类型"]
                    ),
                    "烟叶类别"
                )

            if "致甜成分名称" in current_row_nodes and "甜味持久性" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "持久性",
                        current_row_nodes["甜味持久性"]
                    ),
                    "持久性"
                )

            if "致甜成分名称" in current_row_nodes and "致甜成分分析方法" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "分析方法",
                        current_row_nodes["致甜成分分析方法"]
                    ),
                    "分析方法"
                )

            if "致甜成分名称" in current_row_nodes and "致甜成分实验方法" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "实验方法",
                        current_row_nodes["致甜成分实验方法"]
                    ),
                    "实验方法"
                )

            if "致甜成分名称" in current_row_nodes and "致甜成分添加量及口感" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "添加量及口感",
                        current_row_nodes["致甜成分添加量及口感"]
                    ),
                    "添加量及口感"
                )

            if "致甜成分名称" in current_row_nodes and "致甜成分烟丝中的含量及迁移率" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "含量及迁移率",
                        current_row_nodes["致甜成分烟丝中的含量及迁移率"]
                    ),
                    "含量及迁移率"
                )

            if "致甜成分名称" in current_row_nodes and "致甜成分的优点" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "优点",
                        current_row_nodes["致甜成分的优点"]
                    ),
                    "优点"
                )

            if "致甜成分名称" in current_row_nodes and "致甜成分的缺点" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "缺点",
                        current_row_nodes["致甜成分的缺点"]
                    ),
                    "缺点"
                )

            if "致甜成分名称" in current_row_nodes and "致甜成分的添加部位" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "添加部位",
                        current_row_nodes["致甜成分的添加部位"]
                    ),
                    "添加部位"
                )

            if "致甜成分名称" in current_row_nodes and "致甜成分的香气" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "香气",
                        current_row_nodes["致甜成分的香气"]
                    ),
                    "香气"
                )
            if "致甜成分名称" in current_row_nodes and "致甜成分相对甜度" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "相对甜度",
                        current_row_nodes["致甜成分相对甜度"]
                    ),
                    "相对甜度"
                )

            if "致甜成分名称" in current_row_nodes and "评吸方法" in current_row_nodes:
                graph.merge(
                    Relationship(
                        current_row_nodes["致甜成分名称"],
                        "评吸方式",
                        current_row_nodes["评吸方法"]
                    ),
                    "评吸方式"
                )


    # -------------------------- 5. 导入结果统计 --------------------------
    end_time = time.time()
    total_time = end_time - start_time

    # 统计各标签（列）的节点数量
    print(f"\n📈 导入结果统计：")
    print(f"   • 总耗时：{total_time:.2f} 秒")
    print(f"   • 总创建节点数：{node_count}")
    print(f"   • 各列（标签）节点数量：")

    for col_name in data.columns:
        node_label = col_name
        # 查询该标签的节点总数
        label_count = graph.run(f"MATCH (n:{node_label}) RETURN count(n)").evaluate() or 0
        print(f"     - {col_name}（标签：{node_label}）：{label_count} 个节点")

    print(f"\n🎉 数据导入完成！可在Neo4j Browser查看（http://{neo4j_uri.split('//')[1].split(':')[0]}:7474）")


if __name__ == "__main__":
    # -------------------------- 配置信息（请根据实际情况修改） --------------------------
    CSV_PATH = r"卷烟甜味数据.csv"  # 你的CSV路径
    NEO4J_URI = "neo4j://localhost:7687"  # 用bolt://协议，兼容性更好
    NEO4J_USER = "neo4j"  # 用户名（默认neo4j）
    NEO4J_PASSWORD = "053116wj"  # 你的密码

    # 执行导入
    import_data_to_neo4j(CSV_PATH, NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD)